"""
模型适配器
实现模型动态切换、资源感知、故障转移
根据系统资源状况自动选择最合适的模型
"""

from typing import Optional, Dict, Any, List, TYPE_CHECKING
import asyncio

# 从共享types文件导入类型定义
from middleware.types import ModelType, BaseModelClientProtocol, CloudLLMClientProtocol
from middleware.cloud_llm_client import CloudLLMClient, get_cloud_llm_manager
from config.model_config import model_config_manager, ModelBackend, ModelSize
from app.core.resource_monitor import (
    get_resource_monitor, get_model_resource_info,
    ResourceMonitor
)
# 导入模型客户端 - 移至函数内部避免循环导入
# 导入高级模型管理器
# 使用延迟导入避免循环依赖
from utils.logger import Logger


from utils.constants import MODEL_SWITCH_THRESHOLD

logger = Logger.get_logger("model_adapter")


class ModelAdapter:
    """模型适配器，管理模型的选择和切换"""
    
    def __init__(self):
        self._resource_monitor = get_resource_monitor()
        self._clients: Dict[ModelType, BaseModelClientProtocol] = {}
        self._current_model: Optional[str] = None
        self._preferred_model_type: ModelType = ModelType.OLLAMA
        self._fallback_enabled: bool = True
        self._model_health_status: Dict[str, bool] = {}
        self._switch_lock = asyncio.Lock()
        self._cloud_manager = get_cloud_llm_manager()
        self._using_cloud: bool = False
        self._cloud_client: Optional[CloudLLMClient] = None
    
    async def initialize(self):
        """初始化模型适配器"""
        # 初始化各个后端的客户端
        await self._initialize_clients()
        # 直接使用配置的默认模型
        default_model = model_config_manager.get_current_config().model_name
        if default_model:
            await self._switch_to_model(ModelBackend.OLLAMA, default_model)
        else:
            await self.switch_to_optimal_model()
        logger.info("模型适配器初始化完成")
    
    async def _initialize_clients(self):
        """初始化各模型客户端"""
        try:
            # 导入ModelClientFactory避免循环依赖
            from middleware.model_client import ModelClientFactory
            # 使用ModelClientFactory创建客户端
            # 初始化Fallback客户端作为最后的备选
            self._clients[ModelType.FALLBACK] = ModelClientFactory.create_client(ModelType.FALLBACK)
            
            # 初始化Ollama客户端（默认使用轻量模型）
            self._clients[ModelType.OLLAMA] = ModelClientFactory.create_client(ModelType.OLLAMA)
            
            # 初始化其他客户端
            try:
                self._clients[ModelType.VLLM] = ModelClientFactory.create_client(ModelType.VLLM)
                logger.info("成功初始化VLLM客户端")
            except Exception as e:
                logger.warning(f"初始化VLLM客户端失败: {str(e)}")
            
            try:
                self._clients[ModelType.LLAMA_CPP] = ModelClientFactory.create_client(ModelType.LLAMA_CPP)
                logger.info("成功初始化LLAMA_CPP客户端")
            except Exception as e:
                logger.warning(f"初始化LLAMA_CPP客户端失败: {str(e)}")
        except Exception as e:
            logger.error(f"初始化模型客户端时出错: {str(e)}")
        
        # 初始化健康状态
        for model_type in self._clients:
            self._model_health_status[model_type.value] = True
    
    async def get_client(self) -> BaseModelClientProtocol:
        """获取当前活跃的模型客户端"""
        # 检查是否需要切换模型
        if await self._resource_monitor.should_switch_model():
            await self.switch_to_optimal_model()
        
        # 如果正在使用云端模型，直接返回云端客户端
        if self._using_cloud and self._cloud_client:
            return self._cloud_client
        
        # 根据健康状态和资源选择客户端
        client = await self._select_healthy_client()
        return client
    
    async def _select_healthy_client(self) -> BaseModelClientProtocol:
        """选择一个健康的客户端"""
        # 1. 优先尝试使用云端模型 - 强制使用，不再检查本地模型
        try:
            cloud_client = await self._get_cloud_client()
            if cloud_client:
                logger.info("✓ 使用云端模型客户端")
                print("✓ 使用云端模型客户端")
                return cloud_client
            else:
                logger.info("✗ 云端模型不可用，检查配置是否正确")
                print("✗ 云端模型不可用，检查配置是否正确")
        except Exception as e:
            logger.error(f"✗ 云端模型初始化错误: {str(e)}")
            print(f"✗ 云端模型初始化错误: {str(e)}")
            
        # 2. 仅在云端模型完全不可用时才尝试本地模型
        # 优先选择首选模型类型
        if self._preferred_model_type in self._clients and \
           self._model_health_status.get(self._preferred_model_type.value, True):
            logger.info(f"使用本地首选模型: {self._preferred_model_type}")
            print(f"使用本地首选模型: {self._preferred_model_type}")
            return self._clients[self._preferred_model_type]
        
        # 尝试其他健康的客户端
        for model_type, client in self._clients.items():
            if model_type != ModelType.FALLBACK and \
               self._model_health_status.get(model_type.value, True):
                logger.info(f"使用本地备用模型: {model_type}")
                print(f"使用本地备用模型: {model_type}")
                return client
        
        # 所有都不可用，返回降级客户端
        logger.warning("所有模型后端均不可用，使用降级方案")
        print("使用降级客户端 - 这意味着所有模型都不可用！")
        return self._clients[ModelType.FALLBACK]
    
    async def switch_to_optimal_model(self):
        """切换到最优模型"""
        async with self._switch_lock:
            current_memory = await self._resource_monitor.get_memory_usage()
            
            # 如果内存使用率过高，尝试使用云端模型
            if current_memory >= MODEL_SWITCH_THRESHOLD:
                cloud_client = await self._cloud_manager.get_fallback_client()
                if cloud_client:
                    logger.info(f"内存使用率 {current_memory}% 过高，切换到云端LLM服务")
                    self._using_cloud = True
                    self._cloud_client = cloud_client
                    self._preferred_model_type = ModelType.FALLBACK
                    return
            
            # 内存使用率可接受，使用本地模型
            self._using_cloud = False
            
            # 直接使用资源监控器获取推荐的模型大小
            recommended_size = await self._resource_monitor.get_recommended_model_size()
            logger.info(f"当前内存使用率: {current_memory}%, 推荐模型大小: {recommended_size}")
            
            # 根据推荐大小选择模型
            if recommended_size == "small":
                # 内存紧张，使用小型量化模型
                await self._switch_to_model(ModelBackend.OLLAMA, "phi3:3.8b-instruct-q4_K_M")
            elif recommended_size == "medium":
                # 内存适中，使用中型量化模型
                await self._switch_to_model(ModelBackend.OLLAMA, "qwen2:7b-instruct-q4_K_M")
            else:
                # 内存充足，可使用较大模型
                await self._switch_to_model(ModelBackend.OLLAMA, "llama3")
    
    async def _switch_to_model(self, backend: ModelBackend, model_name: str):
        """切换到指定模型"""
        if self._current_model == model_name:
            return
        
        logger.info(f"切换模型: {self._current_model or 'None'} -> {model_name} (backend: {backend.value})")
        
        # 检查资源使用情况，确保模型可以加载
        current_memory = await self._resource_monitor.get_memory_usage()
        resource_info = get_model_resource_info(model_name)
        
        # 如果内存不足且有资源信息，尝试选择替代模型
        if resource_info and current_memory > 75:  # 内存使用率超过75%时考虑替代模型
            memory_required = resource_info.get("memory_mb", 0)
            available_memory = await self._resource_monitor.get_available_memory_mb()
            
            if memory_required > available_memory * 0.9:  # 如果模型需要的内存超过可用内存的90%
                logger.warning(f"模型 {model_name} 资源占用过高，尝试使用更小的模型")
                # 使用小型模型作为替代
                model_name = "phi3:3.8b-instruct-q4_K_M"
        
        # 更新Ollama配置
        if backend == ModelBackend.OLLAMA:
            try:
                from config.model_config import OllamaConfig
                config = OllamaConfig()
                config.backend_config["model"] = model_name
                config.name = f"ollama-{model_name.split(':')[0]}"
                
                # 注册新配置
                model_config_manager.register_config(config)
                
                # 更新客户端
                from middleware.model_client import ModelClientFactory
                self._clients[ModelType.OLLAMA] = ModelClientFactory.create_client(ModelType.OLLAMA, config=config)
                self._preferred_model_type = ModelType.OLLAMA
                logger.info(f"成功更新Ollama客户端配置为模型: {model_name}")
            except Exception as e:
                logger.error(f"更新Ollama配置失败: {str(e)}")
                # 即使配置更新失败，也要记录当前模型名称
                self._current_model = model_name
        
        # 更新当前模型
        self._current_model = model_name
        
        # 记录模型内存占用
        if resource_info:
            await self._resource_monitor.record_model_memory(
                model_name, resource_info["memory_mb"]
            )
    
    async def handle_model_error(self, error: Exception, model_type: ModelType):
        """处理模型错误，可能触发切换"""
        logger.error(f"模型错误 [{model_type.value}]: {str(error)}")
        
        # 标记模型为不健康
        self._model_health_status[model_type.value] = False
        
        # 如果是资源错误，尝试切换到更小的模型
        if "memory" in str(error).lower() or "oom" in str(error).lower():
            logger.warning("检测到内存错误，尝试切换到更小的模型")
            await self._switch_to_model(ModelBackend.OLLAMA, "phi3:3.8b-instruct-q4_K_M")
        else:
            # 其他错误情况，使用最优模型选择
            await self.switch_to_optimal_model()
        
        # 启动健康检查定时任务
        asyncio.create_task(self._scheduled_health_check(model_type))
    
    async def _scheduled_health_check(self, model_type: ModelType):
        """定时检查模型健康状态"""
        for _ in range(5):  # 最多检查5次
            await asyncio.sleep(30)  # 每30秒检查一次
            
            try:
                client = self._clients.get(model_type)
                if client and await client.health_check():
                    self._model_health_status[model_type.value] = True
                    logger.info(f"模型[{model_type.value}]健康检查通过，恢复使用")
                    return
            except Exception as e:
                logger.error(f"模型健康检查失败: {str(e)}")
    
    async def generate(self, prompt: str, **kwargs) -> Optional[str]:
        """生成文本，自动处理模型选择和故障转移"""
        # 添加超时参数默认值
        if 'timeout' not in kwargs:
            kwargs['timeout'] = 30.0
            
        for attempt in range(3):  # 最多尝试3次
            client = await self.get_client()
            
            try:
                # 如果是云客户端，直接调用generate方法
                if self._using_cloud and isinstance(client, CloudLLMClient):
                    # 云端模型调用添加超时
                    try:
                        result = await asyncio.wait_for(
                            client.generate(prompt, **kwargs), 
                            timeout=kwargs.get('timeout', 30.0)
                        )
                    except asyncio.TimeoutError:
                        logger.error(f"云端模型调用超时 (尝试 {attempt + 1}/3)")
                        # 减少超时时间，准备重试
                        kwargs['timeout'] = max(10.0, kwargs['timeout'] * 0.8)
                        continue
                else:
                    # 本地客户端使用safe_generate或generate
                    try:
                        if hasattr(client, 'safe_generate'):
                            result = await client.safe_generate(prompt, **kwargs)
                        else:
                            # 降级使用generate方法
                            result = await client.generate(prompt, **kwargs)
                    except Exception as e:
                        logger.warning(f"客户端生成失败: {str(e)}，尝试直接使用generate方法")
                        result = await client.generate(prompt, **kwargs)
                
                if result:
                    return result
                
                # 客户端返回None，可能是故障
                if not self._using_cloud:
                    await self.handle_model_error(
                        Exception("模型返回空结果"), 
                        client.model_type
                    )
                
            except Exception as e:
                # 本地模型错误处理
                if not self._using_cloud:
                    await self.handle_model_error(e, client.model_type)
                
                # 尝试使用云端降级
                cloud_client = await self._get_cloud_client()
                if cloud_client:
                    try:
                        logger.warning(f"尝试使用云端LLM降级: {str(e)}")
                        # 添加超时参数，设置为30秒
                        kwargs_with_timeout = {**kwargs, "timeout": 30.0}
                        # 使用wait_for确保超时控制
                        result = await asyncio.wait_for(
                            cloud_client.generate(prompt, **kwargs_with_timeout),
                            timeout=30.0
                        )
                        return result
                    except Exception:
                        pass
                
            # 重试前等待，使用指数退避策略
            wait_time = 0.5 * (2 ** attempt)
            logger.info(f"第 {attempt + 1} 次尝试失败，{wait_time:.2f}秒后重试")
            await asyncio.sleep(wait_time)
        
        # 所有尝试都失败，最后尝试云端降级
        logger.error("所有模型尝试均失败，最后尝试云端降级")
        cloud_client = await self._get_cloud_client()
        if cloud_client:
            try:
                # 添加超时参数，设置为30秒
                kwargs_with_timeout = {**kwargs, "timeout": 30.0}
                # 最终尝试使用更短的超时时间
                result = await asyncio.wait_for(
                    cloud_client.generate(prompt, **kwargs_with_timeout),
                    timeout=20.0
                )
                return result
            except Exception as e:
                logger.error(f"云端LLM降级失败: {str(e)}")
        
        # 所有方案都失败，使用本地降级方案
        return await self._clients[ModelType.FALLBACK].generate(prompt, **kwargs)
    
    async def build_and_generate_qa(self, query: str, contexts: List[str], **kwargs) -> Optional[str]:
        """构建并生成问答内容"""
        client = await self.get_client()
        
        # 为云客户端构建专用的QA prompt
        if self._using_cloud and isinstance(client, CloudLLMClient):
            # 为云模型构建简洁的QA prompt
            context_str = '\n'.join(contexts)
            prompt = f"请基于以下上下文回答问题：\n\n上下文：\n{context_str}\n\n问题：{query}\n\n请提供精确的回答。"
            # 添加超时参数，设置为30秒
            kwargs_with_timeout = {**kwargs, "timeout": 30.0}
            return await self.generate(prompt, **kwargs_with_timeout)
        else:
            # 使用本地客户端的构建方法
            prompt = client.build_qa_prompt(query, contexts)
            return await self.generate(prompt, **kwargs)
    
    async def _get_cloud_client(self) -> Optional[CloudLLMClient]:
        """获取可用的云端LLM客户端 - 增强版，添加详细日志"""
        try:
            print("尝试创建云端客户端...")
            # 检查现有客户端
            if self._cloud_client and self._using_cloud:
                print("使用现有的云端客户端")
                return self._cloud_client
            
            # 尝试获取新的云端客户端
            print("尝试获取新的云端客户端...")
            client = await self._cloud_manager.get_fallback_client()
            if client:
                print("云端客户端获取成功，设置为活动状态")
                self._cloud_client = client
                self._using_cloud = True
                self._preferred_model_type = ModelType.FALLBACK
                return client
            else:
                print("云端客户端获取失败，返回None")
        except Exception as e:
            print(f"获取云端客户端异常: {str(e)}")
            logger.error(f"获取云端客户端失败: {str(e)}")
        
        return None
    
    def is_using_cloud(self) -> bool:
        """检查是否正在使用云端模型"""
        return self._using_cloud
    
    def should_use_cloud_service(self) -> bool:
        """检查是否应该使用云端服务
        
        根据系统资源状况决定是否应该使用云端LLM服务
        注意：此方法使用同步方式获取内存使用情况，避免事件循环问题
        """
        try:
            # 尝试导入psutil库
            import psutil
            
            # 1. 检查是否强制使用云端
            from config.model_config import model_config
            if model_config.get("force_cloud", False):
                logger.info("强制使用云端服务")
                return True
            
            # 2. 检查本地是否有可用资源
            # 检查CPU使用率
            cpu_percent = psutil.cpu_percent(interval=1)
            logger.info(f"当前CPU使用率: {cpu_percent}%")
            
            # 检查内存使用情况
            memory = psutil.virtual_memory()
            memory_percent = memory.percent
            logger.info(f"当前内存使用率: {memory_percent}%")
            
            # 3. 根据资源使用情况决定是否使用云端
            # 如果CPU或内存使用率超过阈值，使用云端服务
            cpu_threshold = 80
            memory_threshold = MODEL_SWITCH_THRESHOLD
            
            if cpu_percent > cpu_threshold or memory_percent > memory_threshold:
                logger.info(f"资源使用过高，使用云端服务: CPU={cpu_percent}%, Memory={memory_percent}%")
                return True
            
            # 4. 检查当前是否已经在使用云端服务
            if self._using_cloud:
                return True
            
            # 5. 检查本地模型健康状态
            if not self._model_health_status.get(self._preferred_model_type.value, False):
                # 如果配置了云端模型，尝试使用
                if hasattr(self, '_clients'):
                    logger.info("本地模型不可用，尝试使用云端服务")
                    return True
            
            logger.info("本地资源充足，使用本地服务")
            return False
        except ImportError:
            logger.warning("psutil库未安装，无法检测系统资源，默认使用本地服务")
            return False
        except Exception as e:
            logger.error(f"检查是否使用云端服务时出错: {str(e)}")
            # 出错时默认使用本地服务
            return False


# 创建全局模型适配器实例
model_adapter = ModelAdapter()


def get_model_adapter() -> ModelAdapter:
    """获取模型适配器实例"""
    # 动态导入避免循环导入
    from middleware.model_client import ModelType
    return model_adapter


async def init_model_adapter():
    """初始化模型适配器"""
    adapter = get_model_adapter()
    await adapter.initialize()
    
    # 初始化云端LLM服务
    from middleware.cloud_llm_client import init_cloud_llm
    await init_cloud_llm()
    
    logger.info("模型适配器初始化完成，云端LLM服务就绪")

# 更新常量文件中的模型切换阈值
import utils.constants
utils.constants.MODEL_SWITCH_THRESHOLD = 85  # 内存使用率超过85%时切换模型